import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
from sklearn.metrics import confusion_matrix
import re

# 设置中文字体
font_path = '/Users/liuyuzhuo/Library/Fonts/SimHei.ttf'  # 请确保此路径正确
font_prop = FontProperties(fname=font_path)
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

def extract_confusion_matrix_from_report():
    """从分类报告中提取混淆矩阵数据"""
    report_path = "output/face_flower_classifier/test_results/classification_report.txt"
    misclassified_path = "output/face_flower_classifier/test_results/all_misclassified_images.txt"
    
    # 从分类报告中提取支持数量
    with open(report_path, 'r') as f:
        report_lines = f.readlines()
    
    # 解析分类报告以获取总支持数
    supports = {}
    for line in report_lines:
        if 'face' in line:
            parts = line.split()
            supports['face'] = int(parts[-1])
        elif 'flower' in line:
            parts = line.split()
            supports['flower'] = int(parts[-1])
    
    # 从误分类文件中计算误分类数量
    face_misclassified = 0
    flower_misclassified = 0
    
    with open(misclassified_path, 'r') as f:
        misclassified_lines = f.readlines()
    
    for line in misclassified_lines:
        if "实际类别: face, 预测为: flower" in line:
            face_misclassified += 1
        elif "实际类别: flower, 预测为: face" in line:
            flower_misclassified += 1
    
    # 构建混淆矩阵
    face_correct = supports['face'] - face_misclassified
    flower_correct = supports['flower'] - flower_misclassified
    
    cm = np.array([
        [face_correct, face_misclassified],
        [flower_misclassified, flower_correct]
    ])
    
    return cm, supports

def plot_confusion_matrix_zh(cm, class_names, supports):
    """绘制中文混淆矩阵"""
    # 创建图表
    fig, ax = plt.subplots(figsize=(10, 8))
    
    # 计算准确率
    total = np.sum(cm)
    accuracy = np.trace(cm) / total * 100
    
    # 创建标准化的混淆矩阵（按行归一化）
    cm_normalized = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
    
    # 设置颜色映射
    cmap = plt.cm.Blues
    
    # 绘制混淆矩阵
    im = ax.imshow(cm_normalized, interpolation='nearest', cmap=cmap)
    
    # 设置坐标轴标签
    tick_marks = np.arange(len(class_names))
    ax.set_xticks(tick_marks)
    ax.set_yticks(tick_marks)
    ax.set_xticklabels(class_names, fontproperties=font_prop, fontsize=14)
    ax.set_yticklabels(class_names, fontproperties=font_prop, fontsize=14)
    
    # 设置轴标签
    ax.set_ylabel('真实类别', fontproperties=font_prop, fontsize=16, labelpad=10)
    ax.set_xlabel('预测类别', fontproperties=font_prop, fontsize=16, labelpad=10)
    
    # 添加标题，包括准确率
    ax.set_title(f'混淆矩阵 (准确率: {accuracy:.2f}%)', fontproperties=font_prop, fontsize=18)
    
    # 添加文本标注到混淆矩阵单元格
    thresh = cm_normalized.max() / 2.
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            text = "{}\n({:.1f}%)".format(cm[i, j], 100 * cm_normalized[i, j])
            ax.text(j, i, text,
                   ha="center", va="center",
                   color="white" if cm_normalized[i, j] > thresh else "black",
                   fontsize=14)
    
    # 添加颜色条
    fig.colorbar(im)
    
    # 保存图表
    plt.tight_layout()
    save_path = "output/face_flower_classifier/test_results/confusion_matrix_zh.png"
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    print(f"中文混淆矩阵已保存至 {save_path}")

def plot_classification_report_zh():
    """生成中文版的分类报告图"""
    report_path = "output/face_flower_classifier/test_results/classification_report.txt"
    
    with open(report_path, 'r') as f:
        report_lines = f.readlines()
    
    # 解析分类报告
    metrics = []
    for line in report_lines:
        if 'face' in line or 'flower' in line or 'accuracy' in line or 'macro avg' in line or 'weighted avg' in line:
            parts = line.split()
            if len(parts) >= 5:  # 确保行有足够的元素
                if 'face' in line:
                    metrics.append(['人脸'] + parts[1:])
                elif 'flower' in line:
                    metrics.append(['花卉'] + parts[1:])
                elif 'accuracy' in line:
                    metrics.append(['准确率'] + parts[1:])
                elif 'macro avg' in line:
                    metrics.append(['宏平均'] + parts[2:])
                elif 'weighted avg' in line:
                    metrics.append(['加权平均'] + parts[2:])
    
    # 创建图表
    plt.figure(figsize=(12, 6))
    
    # 列名和行名
    col_labels = ['精确率', '召回率', 'F1分数', '支持数']
    row_labels = [m[0] for m in metrics]
    
    # 创建表格数据
    cell_data = []
    for m in metrics:
        if len(m) > 1:
            # 对精确率、召回率和F1分数进行格式化
            row_data = [float(val) for val in m[1:4]]
            # 支持数保持为整数
            if len(m) > 4:
                row_data.append(int(m[4]))
            cell_data.append(row_data)
    
    # 创建表格
    table = plt.table(cellText=[['{:.4f}'.format(cell_data[i][j]) if j < 3 else '{}'.format(int(cell_data[i][j])) 
                                for j in range(len(cell_data[i]))] 
                                for i in range(len(cell_data))],
                     rowLabels=row_labels,
                     colLabels=col_labels,
                     loc='center',
                     cellLoc='center')
    
    # 设置表格样式
    table.auto_set_font_size(False)
    table.set_fontsize(12)
    table.scale(1, 1.5)
    
    # 设置行标签和列标签的字体
    for i, label in enumerate(row_labels):
        table._cells[(i+1, -1)]._text.set_fontproperties(font_prop)
    
    for j, label in enumerate(col_labels):
        table._cells[(0, j)]._text.set_fontproperties(font_prop)
    
    # 设置标题
    plt.title('分类报告', fontproperties=font_prop, fontsize=18)
    plt.axis('off')
    
    # 保存图表
    save_path = "output/face_flower_classifier/test_results/classification_report_zh.png"
    plt.tight_layout()
    plt.savefig(save_path, dpi=150)
    plt.close()
    
    print(f"中文分类报告已保存至 {save_path}")

def main():
    # 获取混淆矩阵数据
    cm, supports = extract_confusion_matrix_from_report()
    
    # 绘制中文混淆矩阵
    class_names = ['人脸', '花卉']
    plot_confusion_matrix_zh(cm, class_names, supports)
    
    # 绘制中文分类报告
    plot_classification_report_zh()
    
    print("所有中文图表生成完成！")

if __name__ == "__main__":
    main() 